#!/usr/bin/python
import sys
from astro.utilities import calc_iontau, readatom, convolve_psf
from astro.pyvpfit import VpfitModel
import matplotlib.pyplot as pl
import numpy as np
import commands

help =     """
right arrow, space:      Next spectrum
left arrow, backspace:   Prev spectrum
"""
atom = readatom('/home/nhmc/installed/vpfit/atom.dat')

def plotspec(f26name, ax):
    """ Plots a spectrum and model using the info in a f26 file.
    """
    vp = VpfitModel()
    vp.readf26(f26name)

    # read spectrum
    name = vp.regions.filename[0]
    fh = open(name)
    fwhm = float(fh.next().split()[-1])
    fh.close()
    data = np.loadtxt(name, skiprows=1, unpack=1)
    s = np.rec.fromarrays(data, names='wa,fl,er,co')

    ax.cla()
    ax.set_autoscale_on(0)
    ax.axhline(0,color='0.7')
    ax.plot(s.wa,s.co,':k')
    ax.plot(s.wa,s.er,color='orange')
    ax.plot(s.wa,s.fl,'b',linestyle='steps-mid')

    waticks1 = []
    waticks2 = []
    sumtau = np.zeros_like(s.wa)
    for l in vp.lines:
        ion = l['name'].replace(' ','')
        if l['berr'] > 0 and l['zerr'] > 0 and l['lognerr'] > 0:
            z,logn,b = l['z'], l['logn'], l['b']
            tau,temp = calc_iontau(s.wa, atom[ion], z+1, logn, b, ticks=1)
            sumtau += tau
            if logn < 13.5:
                waticks1.extend(temp)
            else:
                waticks2.extend(temp)
        else:
            print 'skipping: %s' % l

    fwhmpix = (fwhm / 3.0e5 * s.wa.mean()) / np.median(np.diff(s.wa))
    sumtau = convolve_psf(sumtau,fwhmpix)
    fl = np.exp(-sumtau)
    ax.plot(s.wa,fl,'r')
    if waticks1:  ax.vlines(waticks1, 1.1, 1.25, color='k')
    if waticks2:  ax.vlines(waticks2, 1.1, 1.25, color='r')
    ax.axvline(vp.regions.wmin[0],linestyle='dashed',color='0.7')
    ax.axvline(vp.regions.wmax[0],linestyle='dashed',color='0.7')
    ax.set_xlabel('wavelength')
    ax.set_title(
        f26name +
        '    ndrop: %(ndropped)s iter: %(iterations)s nchisq: %(nchisq)s'
        % vp.stats)
    ax.set_ylim(-0.2, 1.4)
    ax.set_xlim(s.wa.min(),s.wa.max())
    pl.minorticks_on()
    return s,vp  # also ticks?

class Junk(object):
    def __init__(self, names, fig):
        self.i = 0
        self.names = names
        self.n = len(names)
        cids = list(fig.canvas.callbacks.callbacks['key_press_event'])
        for cid in cids:
            fig.canvas.callbacks.disconnect(cid)
        self.fig = fig
        self.ax = fig.axes[0]
        self.connect()

    def update(self):
        name = self.names[self.i]
        spec, vp = plotspec(name, self.ax)
        self.fig.canvas.draw()
        print commands.getoutput('cat %s' % name)

    def on_keypress(self, event):
        if event.key in (' ','right'):
            self.i += 1
            if self.i == self.n : self.i -= 1
            self.update()
        elif event.key in ('backspace','left'):
            self.i -= 1
            if self.i == -1:  self.i = 0
            self.update()
    def connect(self):
        self.fig.canvas.mpl_connect('key_press_event', self.on_keypress)

def main():
    fig = pl.figure(figsize=(12,5))
    ax = pl.axes([0.06,0.12,0.9,0.75])
    junk = Junk(sys.argv[1:], fig)
    junk.update()
    if len(sys.argv[1:]) > 1: print help
    pl.show()

if __name__ == '__main__':
    sys.exit(main())
